import os
import pandas as pd
import mysql.connector
import json
from datetime import datetime, timedelta
import re

def is_group_chat(filename):
    """判断是否为群聊"""
    return '@chatroom' in filename

def extract_name_from_filename(filename):
    """从文件名中提取微信ID和昵称"""
    parts = filename.split(' - ')
    if len(parts) >= 2:
        wx_id = parts[0]
        nickname = ' - '.join(parts[1:])
        # 移除可能的文件扩展名
        nickname = re.sub(r'\.xlsx$', '', nickname)
        return wx_id, nickname
    return None, None

def get_chat_time_bounds(df):
    """获取聊天记录的实际开始和结束时间"""
    if df.empty:
        return None, None
    
    # 按时间排序
    df = df.sort_values('消息时间')
    
    # 获取第一条和最后一条消息的时间
    start_time = df['消息时间'].iloc[0]
    end_time = df['消息时间'].iloc[-1]
    
    return start_time, end_time

def split_by_week(df):
    """按周拆分聊天记录"""
    if df.empty:
        return []
    
    # 确保时间列是datetime类型
    try:
        df['消息时间'] = pd.to_datetime(df['消息时间'], errors='coerce')
        # 删除无效的日期
        df = df.dropna(subset=['消息时间'])
        if df.empty:
            print("警告：转换日期后没有有效数据")
            return []
    except Exception as e:
        print(f"转换日期时出错: {e}")
        return []
    
    # 按时间排序
    df = df.sort_values('消息时间')
    
    # 获取第一条消息的时间
    first_message_time = df['消息时间'].iloc[0]
    
    # 计算该周的开始时间（周一）
    week_start = first_message_time - timedelta(days=first_message_time.weekday())
    week_start = week_start.replace(hour=0, minute=0, second=0, microsecond=0)
    
    # 拆分数据
    weekly_data = []
    current_week_start = week_start
    
    while True:
        # 计算当前周的结束时间
        current_week_end = current_week_start + timedelta(days=6, hours=23, minutes=59, seconds=59)
        
        # 筛选当前周的数据
        mask = (df['消息时间'] >= current_week_start) & (df['消息时间'] <= current_week_end)
        week_df = df[mask]
        
        if week_df.empty:
            break
        
        # 获取此周的实际聊天开始和结束时间
        chat_start_time, chat_end_time = get_chat_time_bounds(week_df)
        
        weekly_data.append({
            'week_start': current_week_start,
            'week_end': current_week_end,
            'chat_start_time': chat_start_time,
            'chat_end_time': chat_end_time,
            'data': week_df
        })
        
        # 移动到下一周
        current_week_start = current_week_end + timedelta(seconds=1)
        
        # 检查是否还有数据
        if df['消息时间'].max() <= current_week_end:
            break
    
    return weekly_data

def format_chat_messages(df):
    """将聊天消息格式化为JSON"""
    messages = []
    
    for _, row in df.iterrows():
        if pd.notna(row['发送人']) and pd.notna(row['消息内容']):
            # 确保消息内容是字符串
            sender = str(row['发送人']) if not isinstance(row['发送人'], str) else row['发送人']
            content = str(row['消息内容']) if not isinstance(row['消息内容'], str) else row['消息内容']
            
            # 处理可能的特殊字符
            try:
                messages.append({
                    sender: content
                })
            except Exception as e:
                print(f"处理消息时出错: {e}")
                # 尝试使用更安全的方式添加
                try:
                    messages.append({
                        "Unknown Sender": "消息内容无法解析"
                    })
                except:
                    pass
    
    try:
        return json.dumps(messages, ensure_ascii=False)
    except Exception as e:
        print(f"JSON序列化出错: {e}")
        return json.dumps([{"error": "聊天记录序列化失败"}], ensure_ascii=False)

def process_excel_file(file_path):
    """处理单个Excel文件"""
    filename = os.path.basename(file_path)
    
    # 判断是否为群聊
    is_group = is_group_chat(filename)
    
    # 提取微信ID和昵称
    wx_id, nickname = extract_name_from_filename(filename)
    
    # 读取Excel文件
    try:
        # 尝试使用不同的引擎读取
        try:
            df = pd.read_excel(file_path, engine='openpyxl')
        except Exception as e:
            print(f"使用openpyxl引擎读取失败: {e}")
            try:
                df = pd.read_excel(file_path, engine='xlrd')
            except Exception as e:
                print(f"使用xlrd引擎读取失败: {e}")
                return []
        
        # 检查必要的列是否存在
        required_columns = ['消息时间', '发送人', '消息内容']
        if not all(col in df.columns for col in required_columns):
            print(f"文件 {filename} 缺少必要的列，跳过处理")
            return []
            
        # 删除空行
        df = df.dropna(how='all')
        
        # 按周拆分数据
        weekly_data = split_by_week(df)
        
        # 准备数据库记录
        records = []
        for week_data in weekly_data:
            week_start = week_data['week_start']
            week_end = week_data['week_end']
            chat_start_time = week_data['chat_start_time']
            chat_end_time = week_data['chat_end_time']
            week_df = week_data['data']
            
            # 格式化聊天消息
            chat_history = format_chat_messages(week_df)
            
            # 创建时间跨度JSON
            time_span = json.dumps({
                'start_time': chat_start_time.strftime('%Y-%m-%d %H:%M:%S'),
                'end_time': week_end.strftime('%Y-%m-%d %H:%M:%S')
            }, ensure_ascii=False)
            
            # 准备记录
            record = {
                'nickname': nickname,  # 微信昵称或群聊名
                'is_group': is_group,  # 是否是群聊
                'customer_nickname': '' if is_group else nickname,  # 客户微信昵称，群聊暂时为空
                'chat_history': chat_history,  # 聊天记录
                'ai_analysis': '',  # AI分析记录暂时为空
                'week_start_time': chat_start_time,  # 改为聊天实际开始时间
                'end_time': chat_end_time,  # 聊天实际结束时间
                'time_span': time_span  # 一周跨度
            }
            
            records.append(record)
        
        return records
    
    except Exception as e:
        print(f"处理文件 {filename} 时出错: {e}")
        return []

def update_database(records):
    """更新数据库中的记录"""
    # 数据库配置
    db_config = {
        'host': os.getenv('MYSQL_HOST', 'localhost'),
        'user': os.getenv('MYSQL_USER', 'root'),
        'password': os.getenv('MYSQL_PASSWORD', 'neo..6555894'),
        'database': os.getenv('MYSQL_DATABASE', 'sales_chat_mvp')
    }
    
    try:
        # 连接数据库
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor(dictionary=True)
        
        # 检查end_time字段是否存在，如果不存在则添加
        try:
            cursor.execute("""
            SELECT COUNT(*)
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = %s
            AND TABLE_NAME = 'chat_history'
            AND COLUMN_NAME = 'end_time'
            """, (db_config['database'],))
            column_exists = cursor.fetchone()['COUNT(*)']
            
            if column_exists == 0:
                cursor.execute("""
                ALTER TABLE chat_history
                ADD COLUMN end_time DATETIME COMMENT 'Actual end time of the chat'
                """)
                print("已添加 'end_time' 字段到 'chat_history' 表")
        except Exception as e:
            print(f"检查或添加end_time字段时出错: {e}")
        
        # 获取所有现有记录的nickname
        cursor.execute("""
        SELECT id, nickname
        FROM chat_history
        """)
        existing_records = cursor.fetchall()
        
        # 创建nickname到id的映射
        nickname_to_ids = {}
        for record in existing_records:
            nickname = record['nickname']
            if nickname not in nickname_to_ids:
                nickname_to_ids[nickname] = []
            nickname_to_ids[nickname].append(record['id'])
        
        # 统计更新和插入的记录数
        updated_count = 0
        inserted_count = 0
        
        # 处理记录
        for record in records:
            nickname = record['nickname']
            
            # 检查是否有现有记录
            if nickname in nickname_to_ids and len(nickname_to_ids[nickname]) > 0:
                # 有现有记录，进行更新
                record_id = nickname_to_ids[nickname].pop(0)  # 获取并移除第一个ID
                
                # 更新记录
                update_sql = """
                UPDATE chat_history 
                SET week_start_time = %s, end_time = %s, time_span = %s
                WHERE id = %s
                """
                
                values = (
                    record['week_start_time'].strftime('%Y-%m-%d %H:%M:%S'),
                    record['end_time'].strftime('%Y-%m-%d %H:%M:%S'),
                    record['time_span'],
                    record_id
                )
                
                try:
                    cursor.execute(update_sql, values)
                    updated_count += 1
                    print(f"更新记录 ID: {record_id}, 昵称: {nickname}")
                except Exception as e:
                    print(f"更新记录 ID: {record_id} 失败: {e}")
            else:
                # 没有现有记录，进行插入
                insert_sql = """
                INSERT INTO chat_history (
                    nickname, is_group, customer_nickname, chat_history, 
                    ai_analysis, week_start_time, time_span, end_time
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                """
                
                values = (
                    record['nickname'],
                    record['is_group'],
                    record['customer_nickname'],
                    record['chat_history'],
                    record['ai_analysis'],
                    record['week_start_time'].strftime('%Y-%m-%d %H:%M:%S'),
                    record['time_span'],
                    record['end_time'].strftime('%Y-%m-%d %H:%M:%S')
                )
                
                try:
                    cursor.execute(insert_sql, values)
                    inserted_count += 1
                    print(f"插入新记录, 昵称: {nickname}")
                except Exception as e:
                    print(f"插入记录失败, 昵称: {nickname}, 错误: {e}")
        
        # 提交事务
        conn.commit()
        print(f"成功更新 {updated_count} 条记录，插入 {inserted_count} 条新记录")
    
    except Exception as e:
        print(f"数据库操作时出错: {e}")
    
    finally:
        if 'conn' in locals() and conn.is_connected():
            cursor.close()
            conn.close()

def main():
    """主函数"""
    # 聊天记录目录
    chat_dir = 'history_chat/historychat'
    
    # 获取所有Excel文件
    excel_files = [f for f in os.listdir(chat_dir) if f.endswith('.xlsx') and not f.startswith('.~')]
    
    all_records = []
    
    # 处理每个文件
    for file in excel_files:
        file_path = os.path.join(chat_dir, file)
        records = process_excel_file(file_path)
        all_records.extend(records)
    
    # 更新数据库
    if all_records:
        update_database(all_records)
    else:
        print("没有找到有效的聊天记录")

if __name__ == "__main__":
    main() 